package com.ass.llm.chat.chatImpl;

import cn.hutool.http.HttpRequest;
import cn.hutool.http.HttpResponse;
import cn.hutool.json.JSONArray;
import cn.hutool.json.JSONObject;
import cn.hutool.json.JSONUtil;
import com.ass.llm.chat.LLMChatAbstract;
import com.ass.llm.constant.LLMEnum;
import com.ass.llm.dto.ChatInput;
import com.ass.llm.entity.LlmMessage;
import com.ass.llm.entity.baiLian.Message;
import com.ass.llm.entity.baiLian.MessageBody;
import com.ass.sse.service.SseService;
import com.ass.sse.vo.MessageVo;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.stereotype.Component;

import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStream;
import java.io.InputStreamReader;
import java.util.ArrayList;
import java.util.List;
import java.util.Objects;
import java.util.UUID;
import java.util.function.BiConsumer;
import java.util.function.Consumer;
import java.util.function.Function;

@Component
public class BaiLianChatLLm extends LLMChatAbstract {

    @Value("${LLM.AliLLM.apiKey}")
    private String apiKey;

    @Value("${LLM.AliLLM.chatApiUrl}")
    private String chatApiUrl;


    @Autowired
    private SseService sseService;


    @Override
    public LLMEnum getLLm() {
        return LLMEnum.BAILIAN_LLM;
    }

    @Override
    public String mtdChat(ChatInput chatInput, BiConsumer<LlmMessage,Boolean> llmConsumer,
                          Function<String,LlmMessage> findMessage) {
        //先判断对话id是否为空，如果为空则视为第一次对话
        List<Message> messages = new ArrayList<>();
        String conversationId;
        boolean isInsert = true;
        if (chatInput.getConversationId() == null || chatInput.getConversationId().isEmpty()) {
            //构造请求体
            messages.add(new Message().setRole("system").setContent(chatInput.getSysContext()));
            messages.add(new Message().setRole("user").setContent(chatInput.getQuestion()));
            conversationId = UUID.randomUUID().toString();
        } else {
            //进行查询获得结果之后在构造请求体
            conversationId = chatInput.getConversationId();
            isInsert = false;
            LlmMessage llmMessage = findMessage.apply(conversationId);
            //历史消息
            messages = JSONUtil.parseArray(llmMessage.getMessage()).toList(Message.class);
            //添加当前消息
            messages.add(new Message().setRole("user").setContent(chatInput.getQuestion()));
        }
        //进行调用大模型返回结果
        String assistantContent = sendLLMPostRequest(
                new MessageBody().setMessages(messages)
                        .setModel(chatInput.getModelName())
                        .setTimeout(60000)
                        .setIsStream(chatInput.getIsStream()), chatInput.getSseId());
        //大模型回答结果
        messages.add(new Message().setRole("assistant").setContent(assistantContent));
        LlmMessage llmMessage = new LlmMessage();
        llmMessage.setPlatform(chatInput.getPlatform());
        llmMessage.setModelSupplier(chatInput.getModelSupplier());
        llmMessage.setModelType(chatInput.getModelType());
        llmMessage.setConversationId(conversationId);
        llmMessage.setMessage(JSONUtil.toJsonStr(messages));
        //存储对话结果和用户提问的问题基于对话id
        llmConsumer.accept(llmMessage,isInsert);
        return assistantContent;
    }

    @Override
    public String odaChat(ChatInput chatInput) {
        List<Message> messages = new ArrayList<>();
        //构造请求体
        messages.add(new Message().setRole("system").setContent(chatInput.getSysContext()));
        messages.add(new Message().setRole("user").setContent(chatInput.getQuestion()));
        //进行调用大模型返回结果
        return sendLLMPostRequest(
                new MessageBody().setMessages(messages)
                        .setModel(chatInput.getModelName())
                        .setTimeout(60000)
                        .setIsStream(chatInput.getIsStream()), chatInput.getSseId());
    }


    //发送请求
    protected String sendLLMPostRequest(MessageBody messagesBody, String sseId) {
        // 2. 构造请求体
        JSONObject requestBody = new JSONObject();
        // 将消息列表添加到请求体
        requestBody.set("messages", messagesBody.getMessages());// 将消息列表添加到请求体
        requestBody.set("model", messagesBody.getModel()); // 指定模型
        requestBody.set("stream", messagesBody.getIsStream()); //是否开启流式

        // 4. 发送HTTP请求
        HttpResponse response = HttpRequest.post(chatApiUrl)
                .header("Authorization", "Bearer " + apiKey)
                .header("Content-Type", "application/json")
                .body(JSONUtil.toJsonStr(requestBody)) // 将JSONObject转为字符串
                .timeout(messagesBody.getTimeout()) // 设置超时5秒
                .execute();

        // 5. 处理响应
        if (response.isOk()) {
            return handleResponse(response, messagesBody.getIsStream(), sseId);
        } else {
            return response.body();
        }
    }


    //处理流式输出和阻塞模式下的输出解析
    protected String handleResponse(HttpResponse response, boolean isStream, String sseId) {
        if (isStream) {
            StringBuilder stringBuilder = new StringBuilder();
            if (response != null) {
                InputStream inputStream = response.bodyStream();
                if (inputStream != null) {
                    try {
                        BufferedReader reader = new BufferedReader(new InputStreamReader(inputStream));
                        String line = "";
                        while ((line = reader.readLine()) != null) {
                            if (line.startsWith("data:") && !line.equals("data: [DONE]")) {
                                String data = line.substring(5);
                                JSONObject bodyObject = JSONUtil.parseObj(data);
                                JSONArray choices = JSONUtil.parseArray(bodyObject.get("choices").toString());
                                JSONObject entries = JSONUtil.parseObj(choices.get(0));
                                JSONObject deltaObject = JSONUtil.parseObj(entries.get("delta"));
                                String finishReason = String.valueOf(entries.get("finish_reason"));
                                String content = deltaObject.get("content").toString();

                                MessageVo messageVo = new MessageVo();

                                if(Objects.equals("stop",finishReason)){
                                    messageVo.setName("messageEnd");
                                }else {
                                    messageVo.setName("message");
                                }
                                messageVo.setData(content);
                                messageVo.setClientId(sseId);
                                Thread.sleep(50);
                                stringBuilder.append(content);
                                sseService.sendEvent(sseId, messageVo);
                            }
                        }

                    } catch (IOException e) {
                        e.printStackTrace();
                        throw new RuntimeException(e);
                    }catch (InterruptedException e1){
                        throw new RuntimeException(e1);
                    }
                }
            }

            //返回值要将一句一句的body拼接成一个字符串
            return stringBuilder.toString();
        } else {
            JSONObject parse = JSONUtil.parseObj(JSONUtil.toJsonPrettyStr(response.body()));
            JSONArray choices = JSONUtil.parseArray(parse.get("choices").toString());
            JSONObject entries = JSONUtil.parseObj(choices.get(0));
            JSONObject messageObject = JSONUtil.parseObj(entries.get("message"));
            return messageObject.get("content").toString();
        }
    }


}
